import torch
import torch.nn as nn
import torch.nn.functional as F


class LabelSmoothing(nn.Module):
    """
    NLL loss with label smoothing.
    """

    def __init__(self, smoothing=0.0, is_focal=False):
        """
        Constructor for the LabelSmoothing module.
        :param smoothing: label smoothing factor
        """
        super(LabelSmoothing, self).__init__()
        self.confidence = 1.0 - smoothing
        self.smoothing = smoothing
        self.is_focal = is_focal

    def forward(self, x, target):
        logprobs = torch.nn.functional.log_softmax(x, dim=-1)

        nll_loss = -logprobs.gather(dim=-1, index=target.unsqueeze(1))
        nll_loss = nll_loss.squeeze(1)
        smooth_loss = -logprobs.mean(dim=-1)
        if self.is_focal:
            logit = torch.exp(-nll_loss)
            loss = (1 - logit) * (self.confidence * nll_loss + self.smoothing * smooth_loss)
        else:
            loss = self.confidence * nll_loss + self.smoothing * smooth_loss
        return loss.mean()


class Contrastive_Loss(nn.Module):

    def __init__(self, smoothing_value=0.0, is_focal=False):
        super(Contrastive_Loss, self).__init__()
        if smoothing_value == 0:
            self.loss_fct = nn.CrossEntropyLoss()
        else:
            self.loss_fct = LabelSmoothing(smoothing_value, is_focal)

    def forward(self, x, features, labels):
        part_loss = self.loss_fct(x, labels)
        contrast_loss = self.con_loss(features, labels.view(-1))
        loss = part_loss + contrast_loss
        return loss
    def con_loss(self, features, labels):
        B, _ = features.shape
        features = F.normalize(features)
        cos_matrix = features.mm(features.t())
        pos_label_matrix = torch.stack([labels == labels[i] for i in range(B)]).float()
        neg_label_matrix = 1 - pos_label_matrix
        pos_cos_matrix = 1 - cos_matrix
        neg_cos_matrix = cos_matrix - 0.4
        neg_cos_matrix[neg_cos_matrix < 0] = 0
        loss = (pos_cos_matrix * pos_label_matrix).sum() + (neg_cos_matrix * neg_label_matrix).sum()
        loss /= B
        return loss
